from elasticsearch import Elasticsearch


class ElasticsearchClient:
    def __init__(self, hosts=["http://localhost:9200"], username="elastic", password="elastic", verify_certs=True):
        """
        
        初始化 Elasticsearch 连接
        :param hosts: ES 服务器地址列表，例如 ["http://localhost:9200"]
        :param username: 可选，ES 认证用户名
        :param password: 可选，ES 认证密码
        :param verify_certs: 是否验证 SSL 证书（针对 HTTPS）
        """
        self.es = Elasticsearch(
            hosts=hosts,
            basic_auth=(username, password) if username and password else None,
            verify_certs=verify_certs
        )

    def create_index(self, index_name, settings=None, mappings=None):
        """
        创建索引
        :param index_name: 索引名称
        :param settings: 索引设置（可选）
        :param mappings: 字段映射（可选）
        """
        if self.es.indices.exists(index=index_name):
            print(f"索引 '{index_name}' 已存在")
            return

        body = {}
        if settings:
            body["settings"] = settings
        if mappings:
            body["mappings"] = mappings

        self.es.indices.create(index=index_name, body=body)
        print(f"索引 '{index_name}' 创建成功")

    def insert_document(self, index_name, document):
        """
        插入或更新文档
        :param index_name: 索引名称
        :param document: 文档内容（字典格式）
        """
        res = self.es.index(index=index_name, document=document)
        return res

    def get_document(self, index_name, doc_id):
        """
        获取文档
        :param index_name: 索引名称
        :param doc_id: 文档 ID
        """
        try:
            res = self.es.get(index=index_name, id=doc_id)
            return res["_source"]
        except Exception as e:
            print(f"获取文档失败: {e}")
            return None

    def search(self, index_name, query):
        """
        搜索文档
        :param index_name: 索引名称
        :param query: 查询语句（DSL 结构）
        """
        res = self.es.search(index=index_name, body=query)
        return res["hits"]["hits"]

    def delete_document(self, index_name, doc_id):
        """
        删除文档
        :param index_name: 索引名称
        :param doc_id: 文档 ID
        """
        try:
            res = self.es.delete(index=index_name, id=doc_id)
            return res
        except Exception as e:
            print(f"删除文档失败: {e}")
            return None

    def delete_index(self, index_name):
        """
        删除索引
        :param index_name: 索引名称
        """
        if self.es.indices.exists(index=index_name):
            self.es.indices.delete(index=index_name)
            print(f"索引 '{index_name}' 删除成功")
        else:
            print(f"索引 '{index_name}' 不存在")

    def update_document(self, index: str, doc_id: str, doc: dict):
        try:
            result = self.es.update(
                index=index,
                id=doc_id,
                body={"doc": doc},
                refresh=True
            )
            return result
        except Exception as e:
            print(f"更新文档失败: {e}")
            return None


# 示例用法
if __name__ == "__main__":
    es_client = ElasticsearchClient(hosts=["http://localhost:9200"], username="elastic", password="elastic")

    # 创建索引
    index_name = "test_index"
    mappings = {
        "properties": {
            "title": {"type": "text"},
            "content": {"type": "text"},
            "timestamp": {"type": "date"}
        }
    }
    es_client.create_index(index_name, mappings=mappings)

    # 插入文档
    doc = {"title": "Elasticsearch 教程", "content": "学习如何使用 Elasticsearch", "timestamp": "2024-03-10T12:00:00"}
    es_client.insert_document(index_name, doc_id=1, document=doc)

    # 查询文档
    query = {"query": {"match": {"content": "Elasticsearch"}}}
    print("查询结果:", es_client.search(index_name, query))

    # 获取文档
    print("获取文档:", es_client.get_document(index_name, 1))

    # 删除文档
    es_client.delete_document(index_name, 1)

    # 删除索引
    es_client.delete_index(index_name)